{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agent Executor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "var userHomeDir = System.getProperty(\"user.home\");\n",
    "var localRespoUrl = \"file://\" + userHomeDir + \"/.m2/repository/\";\n",
    "var langchain4jVersion = \"1.0.1\";\n",
    "var langchain4jbeta = \"1.0.1-beta6\";\n",
    "var langgraph4jVersion = \"1.6-SNAPSHOT\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remove installed package from Jupiter cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash \n",
    "rm -rf \\{userHomeDir}/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/bsc/langgraph4j"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "add local maven repo and install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%dependency /add-repo local \\{localRespoUrl} release|never snapshot|always\n",
    "// %dependency /list-repos\n",
    "%dependency /add org.slf4j:slf4j-jdk14:2.0.9\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-core:\\{langgraph4jVersion}\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-langchain4j:\\{langgraph4jVersion}\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-agent-executor:\\{langgraph4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j:\\{langchain4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j-open-ai:\\{langchain4jVersion}\n",
    "\n",
    "%dependency /resolve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Initialize Logger**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "try( var file = new java.io.FileInputStream(\"./logging.properties\")) {\n",
    "    java.util.logging.LogManager.getLogManager().readConfiguration( file );\n",
    "}\n",
    "\n",
    "var log = org.slf4j.LoggerFactory.getLogger(\"AgentExecutor\");\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dev.langchain4j.agent.tool.P;\n",
    "import dev.langchain4j.agent.tool.Tool;\n",
    "\n",
    "import java.util.Optional;\n",
    "\n",
    "import static java.lang.String.format;\n",
    "\n",
    "public class TestTool {\n",
    "    private String lastResult;\n",
    "\n",
    "    Optional<String> lastResult() {\n",
    "        return Optional.ofNullable(lastResult);\n",
    "    }\n",
    "\n",
    "    @Tool(\"tool for test AI agent executor\")\n",
    "    String execTest(@P(\"test message\") String message) {\n",
    "\n",
    "        lastResult = format( \"test tool executed: %s\", message);\n",
    "        return lastResult;\n",
    "    }\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import static org.bsc.langgraph4j.StateGraph.END;\n",
    "import static org.bsc.langgraph4j.StateGraph.START;\n",
    "import org.bsc.langgraph4j.CompileConfig;\n",
    "import org.bsc.langgraph4j.RunnableConfig;\n",
    "import org.bsc.langgraph4j.checkpoint.BaseCheckpointSaver;\n",
    "import org.bsc.langgraph4j.checkpoint.MemorySaver;\n",
    "import org.bsc.langgraph4j.state.AgentState;\n",
    "import org.bsc.langgraph4j.serializer.StateSerializer;\n",
    "\n",
    "import org.bsc.langgraph4j.agentexecutor.AgentExecutor;\n",
    "\n",
    "import dev.langchain4j.model.openai.OpenAiChatModel;\n",
    "\n",
    "import java.util.List;\n",
    "import java.util.Map;\n",
    "import java.util.stream.Collectors;\n",
    "\n",
    "\n",
    "\n",
    "var chatModel = OpenAiChatModel.builder()\n",
    "    .apiKey( System.getenv(\"OPENAI_API_KEY\") )\n",
    "    //.modelName( \"gpt-3.5-turbo-0125\" )\n",
    "    .modelName( \"gpt-4o-mini\" )\n",
    "    .logResponses(true)\n",
    "    .maxRetries(2)\n",
    "    .temperature(0.0)\n",
    "    .maxTokens(2000)\n",
    "    .build();\n",
    "\n",
    "var stateGraph = AgentExecutor.builder()\n",
    "        .chatModel(chatModel)\n",
    "        .toolsFromObject( new TestTool() )\n",
    "        .build();\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test 1 \n",
    "Update State replacing the 'input'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "START \n",
      "callAgent \n",
      "STEP: StateSnapshot{node=__START__, state={messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }]}, config=RunnableConfig{ threadId=test1, checkPointId=27429124-b74c-4eaa-b29e-cbf757f12839, nextNode=agent, streamMode=SNAPSHOTS }} \n",
      "executeTools \n",
      "execute: execTest \n",
      "STEP: StateSnapshot{node=agent, state={messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_3kcsKL1AcgUTXPvr835vdKfv\", name = \"execTest\", arguments = \"{\"message\":\"perform test once\"}\" }] }]}, config=RunnableConfig{ threadId=test1, checkPointId=3e2bf343-4a92-4e71-a215-b42dc807a217, nextNode=action, streamMode=SNAPSHOTS }} \n",
      "callAgent \n",
      "STEP: StateSnapshot{node=action, state={messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_3kcsKL1AcgUTXPvr835vdKfv\", name = \"execTest\", arguments = \"{\"message\":\"perform test once\"}\" }] }, ToolExecutionResultMessage { id = \"call_3kcsKL1AcgUTXPvr835vdKfv\" toolName = \"execTest\" text = \"test tool executed: perform test once\" }]}, config=RunnableConfig{ threadId=test1, checkPointId=b5918cd8-12eb-4b58-a3e3-a4565a70bc81, nextNode=agent, streamMode=SNAPSHOTS }} \n",
      "STEP: StateSnapshot{node=agent, state={agent_response=The test has been executed with the message: \"perform test once\". If you need anything else, feel free to ask!, messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_3kcsKL1AcgUTXPvr835vdKfv\", name = \"execTest\", arguments = \"{\"message\":\"perform test once\"}\" }] }, ToolExecutionResultMessage { id = \"call_3kcsKL1AcgUTXPvr835vdKfv\" toolName = \"execTest\" text = \"test tool executed: perform test once\" }]}, config=RunnableConfig{ threadId=test1, checkPointId=3de4b09a-6801-43f0-9b4b-dd8a40fe63f9, nextNode=__END__, streamMode=SNAPSHOTS }} \n",
      "STEP: NodeOutput{node=__END__, state={agent_response=The test has been executed with the message: \"perform test once\". If you need anything else, feel free to ask!, messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_3kcsKL1AcgUTXPvr835vdKfv\", name = \"execTest\", arguments = \"{\"message\":\"perform test once\"}\" }] }, ToolExecutionResultMessage { id = \"call_3kcsKL1AcgUTXPvr835vdKfv\" toolName = \"execTest\" text = \"test tool executed: perform test once\" }]}} \n"
     ]
    }
   ],
   "source": [
    "import dev.langchain4j.data.message.UserMessage;\n",
    "\n",
    "var saver = new MemorySaver();\n",
    "\n",
    "CompileConfig compileConfig = CompileConfig.builder()\n",
    "                .checkpointSaver( saver )\n",
    "                .build();\n",
    "\n",
    "var graph = stateGraph.compile( compileConfig );\n",
    "\n",
    "var config = RunnableConfig.builder()\n",
    "                .threadId(\"test1\")\n",
    "                .build();    \n",
    "                \n",
    "var iterator = graph.streamSnapshots( Map.of( \"messages\", UserMessage.from(\"perform test once\")), config );  \n",
    "\n",
    "for( var step : iterator ) {\n",
    "    log.info( \"STEP: {}\", step );\n",
    "}\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "StateSnapshot{node=agent, state={agent_response=The test has been executed with the message: \"perform test once\". If you need anything else, feel free to ask!, messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_3kcsKL1AcgUTXPvr835vdKfv\", name = \"execTest\", arguments = \"{\"message\":\"perform test once\"}\" }] }, ToolExecutionResultMessage { id = \"call_3kcsKL1AcgUTXPvr835vdKfv\" toolName = \"execTest\" text = \"test tool executed: perform test once\" }]}, config=RunnableConfig{ threadId=test1, checkPointId=3de4b09a-6801-43f0-9b4b-dd8a40fe63f9, nextNode=__END__, streamMode=VALUES }} \n",
      "StateSnapshot{node=action, state={messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_3kcsKL1AcgUTXPvr835vdKfv\", name = \"execTest\", arguments = \"{\"message\":\"perform test once\"}\" }] }, ToolExecutionResultMessage { id = \"call_3kcsKL1AcgUTXPvr835vdKfv\" toolName = \"execTest\" text = \"test tool executed: perform test once\" }]}, config=RunnableConfig{ threadId=test1, checkPointId=b5918cd8-12eb-4b58-a3e3-a4565a70bc81, nextNode=agent, streamMode=VALUES }} \n",
      "StateSnapshot{node=agent, state={messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_3kcsKL1AcgUTXPvr835vdKfv\", name = \"execTest\", arguments = \"{\"message\":\"perform test once\"}\" }] }]}, config=RunnableConfig{ threadId=test1, checkPointId=3e2bf343-4a92-4e71-a215-b42dc807a217, nextNode=action, streamMode=VALUES }} \n",
      "StateSnapshot{node=__START__, state={messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }]}, config=RunnableConfig{ threadId=test1, checkPointId=27429124-b74c-4eaa-b29e-cbf757f12839, nextNode=agent, streamMode=VALUES }} \n"
     ]
    }
   ],
   "source": [
    "import org.bsc.langgraph4j.state.RemoveByHash;\n",
    "\n",
    "var history = graph.getStateHistory(config).stream().collect( Collectors.toList() );\n",
    "\n",
    "for( var snapshot : history ) {\n",
    "    log.info( \"{}\", snapshot );\n",
    "}\n",
    "\n",
    "\n",
    "// var state2 =  history.get(2);\n",
    "\n",
    "// var updatedConfig = graph.updateState( state2.config(), \n",
    "//     Map.of( \"messages\", UserMessage.from(\"perform test twice\")), null);\n",
    "\n",
    "// var iterator = graph.streamSnapshots( null, updatedConfig );  \n",
    "\n",
    "\n",
    "// for( var step : iterator ) {\n",
    "//     log.info( \"STEP:\\n {}\", step );\n",
    "// }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test 2\n",
    "Update State replacing the 'input' using a cloned state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "START \n",
      "callAgent \n",
      "STEP: StateSnapshot{node=__START__, state={messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }]}, config=RunnableConfig{ threadId=test2, checkPointId=00a2b94c-0a05-495c-8621-3bf6e7a88b72, nextNode=agent, streamMode=SNAPSHOTS }} \n",
      "executeTools \n",
      "execute: execTest \n",
      "STEP: StateSnapshot{node=agent, state={messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_gkU71yQ9RWWls6zVVksyrjGc\", name = \"execTest\", arguments = \"{\"message\":\"perform test once\"}\" }] }]}, config=RunnableConfig{ threadId=test2, checkPointId=887b1fbd-2e30-4b64-9577-3a917533cf68, nextNode=action, streamMode=SNAPSHOTS }} \n",
      "callAgent \n",
      "STEP: StateSnapshot{node=action, state={messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_gkU71yQ9RWWls6zVVksyrjGc\", name = \"execTest\", arguments = \"{\"message\":\"perform test once\"}\" }] }, ToolExecutionResultMessage { id = \"call_gkU71yQ9RWWls6zVVksyrjGc\" toolName = \"execTest\" text = \"test tool executed: perform test once\" }]}, config=RunnableConfig{ threadId=test2, checkPointId=520634ab-1658-4b24-b67c-6bfaf758c119, nextNode=agent, streamMode=SNAPSHOTS }} \n",
      "STEP: StateSnapshot{node=agent, state={agent_response=The test has been executed with the message: \"perform test once\"., messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_gkU71yQ9RWWls6zVVksyrjGc\", name = \"execTest\", arguments = \"{\"message\":\"perform test once\"}\" }] }, ToolExecutionResultMessage { id = \"call_gkU71yQ9RWWls6zVVksyrjGc\" toolName = \"execTest\" text = \"test tool executed: perform test once\" }]}, config=RunnableConfig{ threadId=test2, checkPointId=f3e0ffea-a7f2-461c-a9a7-3e08fa64ea1e, nextNode=__END__, streamMode=SNAPSHOTS }} \n",
      "STEP: NodeOutput{node=__END__, state={agent_response=The test has been executed with the message: \"perform test once\"., messages=[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_gkU71yQ9RWWls6zVVksyrjGc\", name = \"execTest\", arguments = \"{\"message\":\"perform test once\"}\" }] }, ToolExecutionResultMessage { id = \"call_gkU71yQ9RWWls6zVVksyrjGc\" toolName = \"execTest\" text = \"test tool executed: perform test once\" }]}} \n"
     ]
    }
   ],
   "source": [
    "\n",
    "var config = RunnableConfig.builder()\n",
    "                .threadId(\"test2\")\n",
    "                .build();    \n",
    "                \n",
    "                \n",
    "var iterator = graph.streamSnapshots( Map.of( \"messages\", UserMessage.from(\"perform test once\") ), config );  \n",
    "\n",
    "for( var step : iterator ) {\n",
    "    log.info( \"STEP: {}\", step );\n",
    "}\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[UserMessage { name = null contents = [TextContent { text = \"perform test once\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_gkU71yQ9RWWls6zVVksyrjGc\", name = \"execTest\", arguments = \"{\"message\":\"perform test once\"}\" }] }]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "var history = graph.getStateHistory(config).stream().collect( Collectors.toList() );\n",
    "\n",
    "var state2 =  history.get(2);\n",
    "\n",
    "var updatedState = new HashMap<String,Object>();\n",
    "updatedState.putAll(state2.state().data());\n",
    "\n",
    "updatedState.put(  \"messages\", UserMessage.from(\"perform test twice\")  );\n",
    "\n",
    "// var updatedConfig = graph.updateState( state2.config(), updatedState );\n",
    "// log.info( \"UPDATED CONFIG: {}\", updatedConfig );\n",
    "\n",
    "// var iterator = graph.streamSnapshots( null, updatedConfig );  \n",
    "\n",
    "// for( var step : iterator ) {\n",
    "//     log.info( \"STEP:\\n {}\", step );\n",
    "// }    \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Java (rjk 2.2.0)",
   "language": "java",
   "name": "rapaio-jupyter-kernel"
  },
  "language_info": {
   "codemirror_mode": "java",
   "file_extension": ".jshell",
   "mimetype": "text/x-java-source",
   "name": "java",
   "nbconvert_exporter": "script",
   "pygments_lexer": "java",
   "version": "22.0.2+9-70"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
